import argparse
import time
import requests
import pandas as pd
from tqdm.auto import tqdm
import urllib3
urllib3.disable_warnings()

def get_cistrome_regpot(cistrome_id: int):
    url = f"https://db3.cistrome.org/cistrome/samples/{cistrome_id}/regpotential"
    headers = {
        "Accept": "application/json",
        "User-Agent": "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36",
        "Referer": "https://cistrome.org/"
    }
    try:
        with requests.Session() as s:
            # 精确限速逻辑
            elapsed = time.time() - getattr(s, 'last_request', 0)
            if elapsed < 1.0:  # 严格60 RPM控制
                time.sleep(1.0 - elapsed)
                
            response = s.get(
                url,
                headers=headers,
                verify=False,  # 启用SSL验证
                timeout=10
            )
            s.last_request = time.time()
            
            response.raise_for_status()
            data = response.json()
            return pd.DataFrame(data['genes'])[['sym', 'val']].assign(Cistrome_id=cistrome_id)
            
    except Exception as e:
        print(f"ID {cistrome_id} failed: {str(e)}")
        return pd.DataFrame()

def main():
    parser = argparse.ArgumentParser(description='Process Cistrome data')
    parser.add_argument('-i', '--input', required=True, 
                      help='Input TSV file path')
    parser.add_argument('-o', '--output', required=True,
                      help='Output CSV file path')
    parser.add_argument('-b', '--batch-id', type=int, default=0,
                      help='Batch ID (0-10, default=0)')
    args = parser.parse_args()

    # Read and process data
    df = pd.read_csv(args.input, sep='\t')
    processed = (
        df.assign(id=lambda x: x.index)
          .assign(batch_id=lambda x: x['id'] % 11)
          .query(f"batch_id == {args.batch_id}")
          [['Cistrome_id', 'External_id', 'Factor', 'Ontology']]
    )
    
    # Process IDs with rate limiting
    results = []
    for cid in tqdm(processed['Cistrome_id'], desc="Processing"):
        results.append(get_cistrome_regpot(cid))
        time.sleep(1/60)  # Rate limit to ~60 requests/minute
    
    # Combine and save results
    final_df = pd.concat(results, ignore_index=True)
    final_df.to_csv(args.output, index=False)

if __name__ == "__main__":
    main()

